"""
09_subgraphs.py - Subgraphs and Modular Architecture
======================================================
Chapter 11: Breaking complex systems into composable pieces.

KEY CONCEPTS:
  ★ Subgraph = Compiled graph used as a node
  ★ State isolation - Each subgraph has its own state
  ★ Composability - Nest graphs inside graphs

Run: python src/09_subgraphs.py
"""
from typing import TypedDict
from langgraph.graph import StateGraph, START, END

class ResearchState(TypedDict):
    topic: str
    findings: str

def search_topic(state: ResearchState) -> dict:
    return {"findings": f"Found 5 papers on '{state['topic']}'"}

research_graph = StateGraph(ResearchState)
research_graph.add_node("search", search_topic)
research_graph.add_edge(START, "search")
research_graph.add_edge("search", END)
research_sub = research_graph.compile()

class WritingState(TypedDict):
    content: str
    draft: str

def write_draft(state: WritingState) -> dict:
    return {"draft": f"Draft article based on: {state['content'][:60]}..."}

writing_graph = StateGraph(WritingState)
writing_graph.add_node("write", write_draft)
writing_graph.add_edge(START, "write")
writing_graph.add_edge("write", END)
writing_sub = writing_graph.compile()

class OrchestratorState(TypedDict):
    topic: str
    research: str
    article: str

def do_research(state: OrchestratorState) -> dict:
    result = research_sub.invoke({"topic": state["topic"]})
    return {"research": result["findings"]}

def do_writing(state: OrchestratorState) -> dict:
    result = writing_sub.invoke({"content": state["research"]})
    return {"article": result["draft"]}

main_graph = StateGraph(OrchestratorState)
main_graph.add_node("research", do_research)
main_graph.add_node("writing", do_writing)
main_graph.add_edge(START, "research")
main_graph.add_edge("research", "writing")
main_graph.add_edge("writing", END)
app = main_graph.compile()

if __name__ == "__main__":
    print("=" * 60)
    print("09: Subgraphs and Modular Architecture")
    print("=" * 60)
    result = app.invoke({"topic": "AI Agents in Healthcare"})
    print(f"\n  Topic: {result['topic']}")
    print(f"  Research: {result['research']}")
    print(f"  Article: {result['article']}")
    print(f"\n{'=' * 60}\n✅ Subgraphs complete!")
